# datasets.py
import torch
from torchvision import datasets, transforms


def get_dataset(dir, name):  # dir为文件目录路径

	if name == 'mnist':
		train_dataset = datasets.MNIST(dir, train=True, download=True, transform=transforms.ToTensor())
		eval_dataset = datasets.MNIST(dir, train=False, transform=transforms.ToTensor())
		
	elif name == 'cifar':
		# transforms.Compose([])将多种图形变换的方法整合到一起
		transform_train = transforms.Compose([
			transforms.RandomCrop(32, padding=4),  # 填充4个像素，并且随机裁剪为32×32大小的图片
			transforms.RandomHorizontalFlip(),  # 按0.5的概率随机水平翻转PIL图像。
			transforms.ToTensor(),
			transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
		])  # 给定(R, G, B)三个通道的均值和标准差，Normalized_image=(image-mean)/std来将tensor正则化。

		transform_test = transforms.Compose([
			transforms.ToTensor(),
			transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
		])
		
		train_dataset = datasets.CIFAR10(dir, train=True, download=True, transform=transform_train)
		eval_dataset = datasets.CIFAR10(dir, train=False, download=True, transform=transform_test)

	return train_dataset, eval_dataset

# transforms.ToTensor()把一个取值范围是[0,255]的PIL.Image或者shape为(H,W,C)的numpy.ndarray()，转换成
# 形状为[C,H,W]，取值范围是[0,1.0]的torch.FloatTensor
# 所提到的各知识点详见"Pytorch中文文档"
